"""
新闻和基本面数据相关的Celery任务
"""
import logging
import datetime
from celery import group, chain
from celery.utils.log import get_task_logger
from sqlalchemy import func

from modules.tasks.celery_app import app
from modules.models.base import get_db
from modules.models.stock import SecurityInfo
from modules.data_collector.news_fundamental_collector import EconomicNews, MacroEconomicData, FinancialStatement
from modules.data_collector import collector

# 配置任务日志
logger = get_task_logger(__name__)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_financial_news(self, source="sina", limit=50):
    """
    采集财经新闻
    
    Args:
        source: 新闻来源，支持"sina"、"eastmoney"、"10jqka"等
        limit: 采集数量限制
    """
    logger.info(f"开始采集{source}财经新闻，数量限制: {limit}")
    try:
        count = collector.collect_market_data(
            "financial_news", 
            source=source,
            limit=limit
        )
        logger.info(f"成功采集{count}条{source}财经新闻")
        return {
            "status": "success",
            "source": source,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集{source}财经新闻失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_stock_notice(self, stock_code, start_date=None, end_date=None):
    """
    采集股票公告
    
    Args:
        stock_code: 股票代码
        start_date: 开始日期，格式：YYYYMMDD
        end_date: 结束日期，格式：YYYYMMDD
    """
    logger.info(f"开始采集股票 {stock_code} 的公告")
    try:
        count = collector.collect_market_data(
            "stock_notice", 
            code=stock_code,
            start_date=start_date,
            end_date=end_date
        )
        logger.info(f"成功采集{count}条股票 {stock_code} 公告")
        return {
            "status": "success",
            "stock_code": stock_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 公告失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='index_data', max_retries=3, retry_backoff=60)
def collect_macro_economic_data(self, indicator_name, region="中国"):
    """
    采集宏观经济数据
    
    Args:
        indicator_name: 指标名称，如"gdp"、"cpi"、"ppi"等
        region: 地区，默认为"中国"
    """
    logger.info(f"开始采集{region}{indicator_name}宏观经济数据")
    try:
        count = collector.collect_market_data(
            "macro_data", 
            indicator=indicator_name,
            region=region
        )
        logger.info(f"成功采集{count}条{region}{indicator_name}宏观经济数据")
        return {
            "status": "success",
            "indicator": indicator_name,
            "region": region,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集{region}{indicator_name}宏观经济数据失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_financial_statement(self, stock_code, report_type="balance", period="年报"):
    """
    采集财务报表数据
    
    Args:
        stock_code: 股票代码
        report_type: 报表类型，可选值：balance(资产负债表)/income(利润表)/cash_flow(现金流量表)
        period: 报表周期，可选值：年报/一季报/中报/三季报
    """
    logger.info(f"开始采集股票 {stock_code} 的{report_type}财务报表数据，周期: {period}")
    try:
        count = collector.collect_market_data(
            "financial_statement", 
            code=stock_code,
            report_type=report_type,
            period=period
        )
        logger.info(f"成功采集{count}条股票 {stock_code} 的{report_type}财务报表数据")
        return {
            "status": "success",
            "stock_code": stock_code,
            "report_type": report_type,
            "period": period,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 的{report_type}财务报表数据失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_valuation_indicators(self, stock_code):
    """
    采集估值指标数据
    
    Args:
        stock_code: 股票代码
    """
    logger.info(f"开始采集股票 {stock_code} 的估值指标数据")
    try:
        count = collector.collect_market_data(
            "valuation", 
            code=stock_code
        )
        logger.info(f"成功采集{count}条股票 {stock_code} 的估值指标数据")
        return {
            "status": "success",
            "stock_code": stock_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 的估值指标数据失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data')
def schedule_news_collection(self, sources=None, limit_per_source=50):
    """
    调度定时新闻采集任务
    
    Args:
        sources: 新闻来源列表，默认为["sina", "eastmoney"]
        limit_per_source: 每个来源采集的新闻数量
    """
    logger.info("开始调度新闻采集任务")
    
    try:
        if sources is None:
            sources = ["sina", "eastmoney"]
        
        # 创建并行任务组
        tasks = []
        for source in sources:
            tasks.append(collect_financial_news.s(source, limit_per_source))
        
        # 执行任务组
        group(tasks).apply_async()
        
        logger.info(f"已调度 {len(sources)} 个新闻源的采集任务")
        return {
            "status": "success",
            "message": f"已调度 {len(sources)} 个新闻源的采集任务",
            "sources": sources
        }
    except Exception as e:
        logger.error(f"调度新闻采集任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='index_data')
def schedule_macro_data_collection(self):
    """调度宏观经济数据采集任务"""
    logger.info("开始调度宏观经济数据采集任务")
    
    try:
        # 定义要采集的指标和地区
        indicators = [
            {"name": "gdp", "region": "中国"},
            {"name": "cpi", "region": "中国"},
            {"name": "ppi", "region": "中国"},
            {"name": "unemployment", "region": "中国"},
            {"name": "gdp", "region": "美国"},
            {"name": "cpi", "region": "美国"},
            {"name": "gdp", "region": "日本"},
            {"name": "cpi", "region": "欧元区"}
        ]
        
        # 创建并行任务组
        tasks = []
        for indicator in indicators:
            tasks.append(collect_macro_economic_data.s(
                indicator["name"], 
                indicator["region"]
            ))
        
        # 执行任务组
        group(tasks).apply_async()
        
        logger.info(f"已调度 {len(indicators)} 个宏观经济指标的采集任务")
        return {
            "status": "success",
            "message": f"已调度 {len(indicators)} 个宏观经济指标的采集任务"
        }
    except Exception as e:
        logger.error(f"调度宏观经济数据采集任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='stock_data')
def schedule_financial_data_collection(self, stock_list=None, max_stocks=50):
    """
    调度财务数据采集任务
    
    Args:
        stock_list: 股票代码列表，默认为None (自动选择)
        max_stocks: 最大采集股票数量
    """
    logger.info("开始调度财务数据采集任务")
    
    try:
        # 如果未提供股票列表，则从数据库获取
        db = next(get_db())
        
        if stock_list is None:
            # 获取市值排名靠前的股票
            stocks = db.query(SecurityInfo).filter(
                SecurityInfo.type == 'stock',
                SecurityInfo.is_active == 1
            ).order_by(
                # 按ID排序，假设ID越小市值越大（实际可能需要按其他字段排序）
                SecurityInfo.id
            ).limit(max_stocks).all()
            
            stock_list = [stock.code for stock in stocks]
        else:
            # 限制股票数量
            stock_list = stock_list[:max_stocks]
        
        db.close()
        
        logger.info(f"准备采集 {len(stock_list)} 只股票的财务数据")
        
        # 创建任务链
        tasks = []
        for stock_code in stock_list:
            # 对每只股票，采集三种财务报表和估值指标
            subtasks = chain(
                collect_financial_statement.s(stock_code, "balance", "年报"),
                collect_financial_statement.s(stock_code, "income", "年报"),
                collect_financial_statement.s(stock_code, "cash_flow", "年报"),
                collect_valuation_indicators.s(stock_code)
            )
            tasks.append(subtasks)
        
        # 分批执行任务
        batch_size = 10
        for i in range(0, len(tasks), batch_size):
            batch_tasks = tasks[i:i+batch_size]
            for task in batch_tasks:
                task.apply_async()
            logger.info(f"已调度 {i+len(batch_tasks)}/{len(tasks)} 只股票的财务数据采集任务")
        
        # 同时采集公告
        for stock_code in stock_list:
            collect_stock_notice.delay(stock_code)
        
        return {
            "status": "success",
            "message": f"已调度 {len(stock_list)} 只股票的财务数据采集任务"
        }
    except Exception as e:
        logger.error(f"调度财务数据采集任务失败: {str(e)}")
        raise


# 测试函数
def test_news_fundamental_tasks():
    """测试新闻和基本面数据任务函数"""
    # 采集财经新闻
    news_result = collect_financial_news.delay("sina", 10)
    print(f"采集新浪财经新闻任务已提交，任务ID: {news_result.id}")
    
    # 等待任务完成
    result = news_result.get(timeout=120)
    print(f"采集新浪财经新闻结果: {result}")
    
    # 采集宏观经济数据
    macro_result = collect_macro_economic_data.delay("gdp", "中国")
    print(f"采集中国GDP数据任务已提交，任务ID: {macro_result.id}")
    
    # 等待任务完成
    result = macro_result.get(timeout=120)
    print(f"采集中国GDP数据结果: {result}")
    
    # 采集单只股票的财务报表
    db = next(get_db())
    stock = db.query(SecurityInfo).filter(
        SecurityInfo.type == 'stock'
    ).first()
    db.close()
    
    if stock:
        statement_result = collect_financial_statement.delay(stock.code, "balance", "年报")
        print(f"采集股票 {stock.code} 资产负债表任务已提交，任务ID: {statement_result.id}")
        
        # 等待任务完成
        result = statement_result.get(timeout=120)
        print(f"采集股票资产负债表结果: {result}")
    else:
        print("没有找到股票数据，无法测试采集任务")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_news_fundamental_tasks()